# 封装了对词向量数据库的调用
# 根据传入的字符串查询前三个结果


# app.py
from flask import Flask, request, jsonify
from flask_cors import CORS  # 处理跨域
import os
from typing import List, Dict
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings

# 配置常量
MODELS_PATH = "embedding/models/m3e-base"  # 根据实际情况修改
PERSIST_PATH = "embedding/database/chroma"

# 初始化Flask应用
app = Flask(__name__)
CORS(app)  # 允许所有域的跨域请求


# 初始化向量数据库（在应用启动时加载）
def get_embedding(model_type: str, path):
	if model_type == 'm3e':
		return HuggingFaceEmbeddings(
			model_name=path,
			# model_kwargs={'device': 'cpu'},  # 使用CPU推理
			# encode_kwargs={'normalize_embeddings': True}  # 标准化向量
		)
	else:
		raise ValueError(f"不支持的嵌入模型类型: {model_type}")


def initialize_vector_db() -> Chroma:
	print("初始化向量数据库...")
	return Chroma(
		persist_directory=PERSIST_PATH,
		embedding_function=get_embedding("m3e", MODELS_PATH)
	)


# 全局数据库实例
vectordb = initialize_vector_db()


# @app.route('/api/query', methods=['POST'])
# def handle_query():
# 	"""
# 	处理查询请求
# 	JSON输入格式: {"query": "你的查询文本"}
# 	"""
# 	# 获取请求数据
# 	print("方法启动")
# 	data = request.get_json()
# 	print(data)
#
# 	if not data or 'query' not in data:
# 		return jsonify({"error": "Missing 'query' parameter"}), 400
#
# 	query_text = data['query']
# 	print(query_text)
#
# 	try:
# 		# 执行相似性搜索
# 		sim_docs = vectordb.similarity_search(query_text, k=3)
# 		# 格式化结果
# 		results = []
# 		for doc in sim_docs:
# 			results.append({
# 				"content": doc.page_content[:200] + '...'  # 截取前200字符
# 			})
#
# 		return jsonify({"results": results})
#
# 	except Exception as e:
# 		return jsonify({"error": str(e)}), 500
@app.route('/api/query', methods=['POST'])
def handle_query():
	try:
		data = request.get_json()
		if not data or 'query' not in data:
			return jsonify({"error": "Missing 'query' parameter"}), 400

		query_text = data['query']
		# 执行相似性搜索
		sim_docs = vectordb.similarity_search(query_text, k=3)
		results = [{"content": doc.page_content[:200] + '...'} for doc in sim_docs]
		return jsonify({"results": results})

	except Exception as e:
		print("Error occurred: ", e)  # 打印详细错误信息
		return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
	app.run(port=5010, debug=True)
